import torch
from model.simplenet import SimpleNet
from utils.datasets import SimpleDatasetClass
# from utils.loss import bounding_box_loss
from utils.utils import freeze_param

lr = 1e-5
train_dataset = SimpleDatasetClass()
train_dataloader = torch.utils.data.DataLoader(train_dataset, batch_size=32, shuffle=True)
print('Dataset loading complete.')
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model = SimpleNet()
model.load_state_dict(torch.load('resnet18.pth'))
freeze_param(model, exclude='classify')
optimizer = torch.optim.SGD(model.parameters(), lr)
loss_function = torch.nn.BCELoss()
model.to(device)
print('Network loading complete.')
train_loss = 0
epoch = 0
while True:
    for step, data in enumerate(train_dataloader, start=0):
        images, labels = data
        preds_train = model(images.to(device))
        train_loss = loss_function(torch.sigmoid(preds_train[:, 0]), labels.to(device))
        optimizer.zero_grad()
        train_loss.backward()
        optimizer.step()
        print(f'\rloss={train_loss.item()}', end='')
    print(f'\nEpoch{epoch} finished. Press y to continue, or press other key to finish')
    epoch += 1
    presskey = input()
    if presskey == 'y' or presskey == 'Y':
        continue
    else:
        break
print('Press y to save weights file, or press other key to finish')
presskey = input()
if presskey == 'y' or presskey == 'Y':
    torch.save(model.state_dict(), 'resnet18.pth')
